import codecs
import re
import sys
import nltk
from collections import OrderedDict


def filter_event(_event):
    segs = re.split(r' <V> | <A1> | <A2> ', _event)
    if len(segs) < 2:
        return False
    for seg in segs:
        if len(seg.split()) > 5:
            return False
    return True


def change_event(_event):
    temp = _event
    if '> <' in temp:
        temp = re.sub(r'<[AV][012]?> <([AV])', r'<\1', temp)
        temp = re.sub(r'<SEP> <SEP>', '<SEP>', temp)
    return temp


def collect_entity(_events):
    _entity = OrderedDict()
    for e in _events:
        e = re.sub(r' <V> \S+', '', e)
        segs = re.split(r' <S> | <A1> | <A2> ', e)
        for seg in segs:
            if ' ' not in seg and '<NE>' not in seg and '<ent>' not in seg:
                pos = nltk.pos_tag([seg])[0][1]
                if 'NN' not in pos:
                    continue
            if seg not in _entity:
                _entity[seg] = 1
            else:
                _entity[seg] += 1
    _entity = [item[0] for item in _entity.items() if item[1] > 1]
    return _entity


event_path = sys.argv[1]
orderd_tgt_event = sys.argv[2]
stand_tgt_event = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.srl2.format.event.replace'
stand_tgt_story = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.srl2.format.new'

story_events = codecs.open(event_path, 'r', encoding='utf8').readlines()
tgt_event = codecs.open(orderd_tgt_event, 'r', encoding='utf8').readlines()
org_event = codecs.open(stand_tgt_event, 'r', encoding='utf8').readlines()
org_sotry = codecs.open(stand_tgt_story, 'r', encoding='utf8').readlines()
assert len(story_events) == len(tgt_event)
assert len(org_event) == len(org_sotry)
event_map = zip(org_event, range(len(org_event)))
event_map = dict(event_map)


result_event = event_path + '.formatevent'
result_entity = event_path + '.coref'
result_tgt = event_path + '.tgt'


story_format_events = []
story_entity = []
story_tgt = []

for i, story in enumerate(story_events):
    # story = story.strip().replace(' <SEP>', '')
    story = story.strip()
    if story.startswith('<A0> '):
        story = story[5:]
    keywords = (' <A0>', ' <V>', ' <A1>', ' <A2>', ' <S>')
    if story.endswith(keywords):
        story = story[:-5]
    story = story.replace('<unk>', '#unk#')
    events = story.split(' <A0> ')

    format_event = filter(filter_event, events)
    format_event = map(change_event, format_event)
    dump_event = []
    for e in format_event:
        if len(dump_event) > 0 and e == dump_event[-1]:
            continue
        else:
            dump_event.append(e)
    format_event = dump_event
    entity = collect_entity(format_event)

    format_event = '<A0> ' + ' <A0> '.join(format_event)
    format_event = format_event.replace('#unk#', '<unk>')
    entity.insert(0, '<ent>0')
    entity = map(lambda x: (x + '@0'), entity)
    entity = ','.join(entity)
    entity = entity.replace('#unk#', '<unk>')

    story_format_events.append(format_event)
    story_entity.append(entity)

    text = tgt_event[i]
    if i == len(story_events) - 1:
        text = text + '\n'
    org_index = event_map[text]
    story_tgt.append(org_sotry[org_index])

with codecs.open(result_event, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(story_format_events))
    fout.write('\n')

with codecs.open(result_entity, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(story_entity))
    fout.write('\n')

with codecs.open(result_tgt, 'w', encoding='utf8') as fout:
    fout.write(''.join(story_tgt))